#include "Protocol.h"
#include "rapidjson/writer.h"
#include "rapidjson/stringbuffer.h"
#include "rapidjson/document.h"
#include <iostream>

// 将请求结构转换为 JSON 字符串
std::string Protocol::SerializeRequest(const Request& request) {
  rapidjson::Document doc;
  doc.SetObject();
  rapidjson::Document::AllocatorType& allocator = doc.GetAllocator();

  // 添加模型名称
  doc.AddMember("Model", rapidjson::StringRef(request.model.c_str()), allocator);

  // 添加消息列表
  rapidjson::Value messages(rapidjson::kArrayType);
  for (const auto& message : request.messages) {
    rapidjson::Value msg(rapidjson::kObjectType);
    msg.AddMember("Role", rapidjson::StringRef(message.role.c_str()), allocator);
    msg.AddMember("Content", rapidjson::StringRef(message.content.c_str()), allocator);
    messages.PushBack(msg, allocator);
  }
  doc.AddMember("Messages", messages, allocator);

  // 添加最大 token 数
  doc.AddMember("Temperature", request.temperature, allocator);
  // 添加最大 token 数
  doc.AddMember("Max_tokens", request.max_tokens, allocator);

  // 将 JSON 转换为字符串
  rapidjson::StringBuffer buffer;
  rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
  doc.Accept(writer);
  std::string json_str = buffer.GetString();

  // 打印请求体
  std::cout << "[Protocol] Serialized Request: " << json_str << std::endl;

  return json_str;
}

// 将 JSON 字符串解析为响应结构
Response Protocol::ParseResponse(const std::string& json) {
  Response response;
  response.error = false;  // 默认设置为 false
  rapidjson::Document doc;
  doc.Parse(json.c_str());

  // 检查是否解析成功
  if (doc.HasParseError()) {
    response.error = true;
    response.error_message = "Failed to parse JSON response.";
    std::cerr << "[Protocol] JSON Parse Error: " << doc.GetParseError() << std::endl;
    return response;
  }

  // 解析响应字段
  if (doc.HasMember("id") && doc["id"].IsString()) {
    response.id = doc["id"].GetString();
  }
  if (doc.HasMember("object") && doc["object"].IsString()) {
    response.object = doc["object"].GetString();
  }
  if (doc.HasMember("created") && doc["created"].IsInt()) {
    response.created = doc["created"].GetInt();
  }
  if (doc.HasMember("model") && doc["model"].IsString()) {
    response.model = doc["model"].GetString();
  }

  // 解析 choices 字段
  if (doc.HasMember("choices") && doc["choices"].IsArray()) {
    const rapidjson::Value& choices = doc["choices"];
    for (const auto& choice : choices.GetArray()) {
      if (choice.HasMember("message") && choice["message"].IsObject()) {
        const rapidjson::Value& message = choice["message"];
        Message msg;
        if (message.HasMember("role") && message["role"].IsString()) {
          msg.role = message["role"].GetString();
        }
        if (message.HasMember("content") && message["content"].IsString()) {
          msg.content = message["content"].GetString();
        }
        if (message.HasMember("reasoning_content") && message["reasoning_content"].IsString()) {
          msg.reasoning_content = message["reasoning_content"].GetString();
        }
        response.choices.push_back(msg);
      }
    }
  }

  // 检查是否存在错误
  if (doc.HasMember("error") && doc["error"].IsObject()) {
    response.error = true;
    const rapidjson::Value& error = doc["error"];
    if (error.HasMember("message") && error["message"].IsString()) {
      response.error_message = error["message"].GetString();
    }
  }

  // 打印解析后的响应
  std::cout << "[Protocol] Parsed Response: " << std::endl;
  std::cout << "  ID: " << response.id << std::endl;
  std::cout << "  Model: " << response.model << std::endl;
  std::cout << "  Error: " << (response.error ? "Yes" : "No") << std::endl;
  if (response.error) {
    std::cout << "  Error Message: " << response.error_message << std::endl;
  }
  for (const auto& choice : response.choices) {
    std::cout << "  Choice: " << choice.role << " - " << choice.reasoning_content << " - " << choice.content << std::endl;
  }

  return response;
}